import torchvision
from torch import nn
import torch
import ClassifyModels.utils as utils

"""
using resnet18 to classify, which pre-trained on IMAGENET.
"""

# 数据集的迭代器，自己指定数据集
train_iter, test_iter = None, None

if __name__ == '__main__':
    model = torchvision.models.resnet18(pretrained=True)
    model.fc = nn.LazyLinear(2)
    devices = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    epoches = 10
    lr = 1e-5
    loss = nn.CrossEntropyLoss(reduction="none")
    params_1xlr = [param for name, param in model.named_parameters()
                   if name not in ["fc.weight", "fc.bias"]]
    sgd_opter = torch.optim.SGD([{'params': params_1xlr},
                                 {'params': model.fc.parameters(),
                                  'lr': lr * 10}],
                                lr=lr, weight_decay=0.001)
    with utils.Timer('Train-Test'):
        for epoch in range(epoches):
            utils.train(epoch, train_iter, model, loss, sgd_opter, devices)
            utils.test(epoch, test_iter, model, loss, sgd_opter, devices)
